/*
 * Copyright 1999-2017 Alibaba Group Holding Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package studio.raptor.sqlparser.dialect.postgresql.parser;

import java.util.List;
import studio.raptor.sqlparser.ast.SQLExpr;
import studio.raptor.sqlparser.ast.SQLName;
import studio.raptor.sqlparser.ast.SQLStatement;
import studio.raptor.sqlparser.ast.expr.SQLCurrentOfCursorExpr;
import studio.raptor.sqlparser.ast.expr.SQLIdentifierExpr;
import studio.raptor.sqlparser.ast.expr.SQLQueryExpr;
import studio.raptor.sqlparser.ast.statement.SQLAlterTableAlterColumn;
import studio.raptor.sqlparser.ast.statement.SQLColumnDefinition;
import studio.raptor.sqlparser.ast.statement.SQLInsertStatement;
import studio.raptor.sqlparser.ast.statement.SQLSelect;
import studio.raptor.sqlparser.ast.statement.SQLTableSource;
import studio.raptor.sqlparser.ast.statement.SQLUpdateStatement;
import studio.raptor.sqlparser.dialect.postgresql.ast.PGWithClause;
import studio.raptor.sqlparser.dialect.postgresql.ast.PGWithQuery;
import studio.raptor.sqlparser.dialect.postgresql.ast.stmt.PGDeleteStatement;
import studio.raptor.sqlparser.dialect.postgresql.ast.stmt.PGInsertStatement;
import studio.raptor.sqlparser.dialect.postgresql.ast.stmt.PGSelectStatement;
import studio.raptor.sqlparser.dialect.postgresql.ast.stmt.PGShowStatement;
import studio.raptor.sqlparser.dialect.postgresql.ast.stmt.PGUpdateStatement;
import studio.raptor.sqlparser.parser.Lexer;
import studio.raptor.sqlparser.parser.ParserException;
import studio.raptor.sqlparser.parser.SQLSelectParser;
import studio.raptor.sqlparser.parser.SQLStatementParser;
import studio.raptor.sqlparser.parser.Token;

public class PGSQLStatementParser extends SQLStatementParser {

  public PGSQLStatementParser(String sql) {
    super(new PGExprParser(sql));
  }

  public PGSQLStatementParser(Lexer lexer) {
    super(new PGExprParser(lexer));
  }

  public PGSelectParser createSQLSelectParser() {
    return new PGSelectParser(this.exprParser);
  }

  public SQLUpdateStatement parseUpdateStatement() {
    accept(Token.UPDATE);

    PGUpdateStatement udpateStatement = new PGUpdateStatement();

    SQLSelectParser selectParser = this.exprParser.createSelectParser();
    SQLTableSource tableSource = selectParser.parseTableSource();
    udpateStatement.setTableSource(tableSource);

    parseUpdateSet(udpateStatement);

    if (lexer.token() == Token.FROM) {
      lexer.nextToken();
      SQLTableSource from = selectParser.parseTableSource();
      udpateStatement.setFrom(from);
    }

    if (lexer.token() == (Token.WHERE)) {
      lexer.nextToken();
      udpateStatement.setWhere(this.exprParser.expr());
    }

    if (lexer.token() == Token.RETURNING) {
      lexer.nextToken();

      for (; ; ) {
        udpateStatement.getReturning().add(this.exprParser.expr());
        if (lexer.token() == Token.COMMA) {
          lexer.nextToken();
          continue;
        }
        break;
      }
    }

    return udpateStatement;
  }

  public PGInsertStatement parseInsert() {
    PGInsertStatement stmt = new PGInsertStatement();

    if (lexer.token() == Token.INSERT) {
      lexer.nextToken();
      accept(Token.INTO);

      SQLName tableName = this.exprParser.name();
      stmt.setTableName(tableName);

      if (lexer.token() == Token.IDENTIFIER) {
        stmt.setAlias(lexer.stringVal());
        lexer.nextToken();
      }

    }

    if (lexer.token() == Token.DEFAULT) {
      lexer.nextToken();
      accept(Token.VALUES);
      stmt.setDefaultValues(true);
    }

    if (lexer.token() == (Token.LPAREN)) {
      lexer.nextToken();
      this.exprParser.exprList(stmt.getColumns(), stmt);
      accept(Token.RPAREN);
    }

    if (lexer.token() == (Token.VALUES)) {
      lexer.nextToken();

      for (; ; ) {
        accept(Token.LPAREN);
        SQLInsertStatement.ValuesClause valuesCaluse = new SQLInsertStatement.ValuesClause();
        this.exprParser.exprList(valuesCaluse.getValues(), valuesCaluse);
        stmt.addValueCause(valuesCaluse);

        accept(Token.RPAREN);
        if (lexer.token() == Token.COMMA) {
          lexer.nextToken();
          continue;
        }
        break;
      }
    } else if (lexer.token() == (Token.SELECT)) {
      SQLQueryExpr queryExpr = (SQLQueryExpr) this.exprParser.expr();
      stmt.setQuery(queryExpr.getSubQuery());
    }

    if (lexer.token() == Token.RETURNING) {
      lexer.nextToken();
      SQLExpr returning = this.exprParser.expr();
      stmt.setReturning(returning);
    }
    return stmt;
  }

  public PGDeleteStatement parseDeleteStatement() {
    lexer.nextToken();
    PGDeleteStatement deleteStatement = new PGDeleteStatement();

    if (lexer.token() == (Token.FROM)) {
      lexer.nextToken();
    }
    if (lexer.token() == (Token.ONLY)) {
      lexer.nextToken();
      deleteStatement.setOnly(true);
    }

    SQLName tableName = exprParser.name();

    deleteStatement.setTableName(tableName);

    if (lexer.token() == Token.AS) {
      accept(Token.AS);
    }
    if (lexer.token() == Token.IDENTIFIER) {
      deleteStatement.setAlias(lexer.stringVal());
      lexer.nextToken();
    }

    if (lexer.token() == Token.USING) {
      lexer.nextToken();
      for (; ; ) {
        SQLName name = this.exprParser.name();
        deleteStatement.getUsing().add(name);
        if (lexer.token() == Token.COMMA) {
          lexer.nextToken();
          continue;
        }
        break;
      }
    }

    if (lexer.token() == (Token.WHERE)) {
      lexer.nextToken();

      if (lexer.token() == Token.CURRENT) {
        lexer.nextToken();
        accept(Token.OF);
        SQLName cursorName = this.exprParser.name();
        SQLExpr where = new SQLCurrentOfCursorExpr(cursorName);
        deleteStatement.setWhere(where);
      } else {
        SQLExpr where = this.exprParser.expr();
        deleteStatement.setWhere(where);
      }
    }

    if (lexer.token() == Token.RETURNING) {
      lexer.nextToken();
      accept(Token.STAR);
      deleteStatement.setReturning(true);
    }

    return deleteStatement;
  }

  public boolean parseStatementListDialect(List<SQLStatement> statementList) {
    if (lexer.token() == Token.WITH) {
      SQLStatement stmt = parseWith();
      statementList.add(stmt);
      return true;
    }

    return false;
  }

  public PGWithClause parseWithClause() {
    lexer.nextToken();

    PGWithClause withClause = new PGWithClause();

    if (lexer.token() == Token.RECURSIVE) {
      lexer.nextToken();
      withClause.setRecursive(true);
    }

    for (; ; ) {
      PGWithQuery withQuery = withQuery();
      withClause.getWithQuery().add(withQuery);
      if (lexer.token() == Token.COMMA) {
        lexer.nextToken();
        continue;
      } else {
        break;
      }
    }
    return withClause;
  }

  private PGWithQuery withQuery() {
    PGWithQuery withQuery = new PGWithQuery();

    if (lexer.token() == Token.LITERAL_ALIAS) {
      withQuery.setName(new SQLIdentifierExpr("\"" + lexer.stringVal()
          + "\""));
    } else {
      withQuery.setName(new SQLIdentifierExpr(lexer.stringVal()));
    }
    lexer.nextToken();

    if (lexer.token() == Token.LPAREN) {
      lexer.nextToken();

      for (; ; ) {
        SQLExpr expr = this.exprParser.expr();
        withQuery.addColumn(expr);
        if (lexer.token() == Token.COMMA) {
          lexer.nextToken();
          continue;
        } else {
          break;
        }
      }

      accept(Token.RPAREN);
    }

    accept(Token.AS);

    if (lexer.token() == Token.LPAREN) {
      lexer.nextToken();

      SQLStatement query;
      if (lexer.token() == Token.SELECT) {
        query = this.parseSelect();
      } else if (lexer.token() == Token.INSERT) {
        query = this.parseInsert();
      } else if (lexer.token() == Token.UPDATE) {
        query = this.parseUpdateStatement();
      } else if (lexer.token() == Token.DELETE) {
        query = this.parseDeleteStatement();
      } else if (lexer.token() == Token.VALUES) {
        query = this.parseSelect();
      } else {
        throw new ParserException("syntax error, support token '" + lexer.token() + "'");
      }
      withQuery.setQuery(query);

      accept(Token.RPAREN);
    }

    return withQuery;
  }

  public PGSelectStatement parseSelect() {
    PGSelectParser selectParser = createSQLSelectParser();
    SQLSelect select = selectParser.select();
    return new PGSelectStatement(select);
  }

  public SQLStatement parseWith() {
    PGWithClause with = this.parseWithClause();
    if (lexer.token() == Token.INSERT) {
      PGInsertStatement stmt = this.parseInsert();
      stmt.setWith(with);
      return stmt;
    }

    if (lexer.token() == Token.SELECT) {
      PGSelectStatement stmt = this.parseSelect();
      stmt.setWith(with);
      return stmt;
    }

    if (lexer.token() == Token.DELETE) {
      PGDeleteStatement stmt = this.parseDeleteStatement();
      stmt.setWith(with);
      return stmt;
    }
    throw new ParserException("TODO");
  }

  protected SQLAlterTableAlterColumn parseAlterColumn() {
    if (lexer.token() == Token.COLUMN) {
      lexer.nextToken();
    }

    SQLColumnDefinition column = this.exprParser.parseColumn();

    SQLAlterTableAlterColumn alterColumn = new SQLAlterTableAlterColumn();
    alterColumn.setColumn(column);

    if (column.getDataType() == null && column.getConstraints().size() == 0) {
      if (lexer.token() == Token.SET) {
        lexer.nextToken();
        if (lexer.token() == Token.NOT) {
          lexer.nextToken();
          accept(Token.NULL);
          alterColumn.setSetNotNull(true);
        } else {
          accept(Token.DEFAULT);
          SQLExpr defaultValue = this.exprParser.expr();
          alterColumn.setSetDefault(defaultValue);
        }
      } else if (lexer.token() == Token.DROP) {
        lexer.nextToken();
        if (lexer.token() == Token.NOT) {
          lexer.nextToken();
          accept(Token.NULL);
          alterColumn.setDropNotNull(true);
        } else {
          accept(Token.DEFAULT);
          alterColumn.setDropDefault(true);
        }
      }
    }
    return alterColumn;
  }

  public SQLStatement parseShow() {
    accept(Token.SHOW);
    PGShowStatement stmt = new PGShowStatement();
    stmt.setExpr(this.exprParser.expr());
    return stmt;
  }
}
